import warnings

import jieba
from langchain.retrievers import BM25Retriever
from langchain.schema import Document

from utils.utils import timer

warnings.filterwarnings("ignore")

# 可选：加载自定义词典（提升分词质量）
# jieba.load_userdict("path/to/custom_dict.txt")


def _tokenize(text):
    """高效中文分词"""
    return " ".join(jieba.lcut(text))


class BM25:
    """基于 BM25 的中文文档检索器，优化性能与内存。"""

    @timer
    def __init__(self, documents):
        self._original_docs = []  # 存原始 Document（未分词）
        tokenized_documents = []  # 临时用于构建 retriever

        for idx, line in enumerate(documents):
            line = line.strip()
            if len(line) < 5:
                continue
            # 提取原始内容（tab 分割第一部分）
            content = line.split("\t")[0]
            # 构建原始文档
            self._original_docs.append(
                Document(page_content=content, metadata={"id": idx})
            )
            # 构建用于 BM25 的分词文档
            tokens = _tokenize(line)
            tokenized_documents.append(
                Document(page_content=tokens, metadata={"id": idx})
            )

        # 一次性构建 retriever
        self._retriever = BM25Retriever.from_documents(tokenized_documents)

    @timer
    def get_bm25_topk(self, query, topk):
        """
        检索 topk 相关文档（性能优化版）。
        """
        tokenized_query = _tokenize(query)
        # 直接传 k 参数，避免修改 retriever 状态
        retrieved_docs = self._retriever.get_relevant_documents(tokenized_query, k=topk)
        # 映射回原始文档
        return [self._original_docs[doc.metadata["id"]] for doc in retrieved_docs]


if __name__ == "__main__":
    from pdf_parse import DataProcess

    # 初始化 PDF 解析器
    dataprocess = DataProcess(pdf_path="./data/train_a.pdf")

    # 多粒度文本解析
    print("🔍 开始多粒度块解析...")
    dataprocess.parse_blocks(max_seq=1024)
    dataprocess.parse_blocks(max_seq=512)
    print(f"  块解析后文档数: {len(dataprocess.data)}")

    print("\n📖 开始全页解析...")
    dataprocess.parse_all_page(max_seq=256)
    dataprocess.parse_all_page(max_seq=512)
    print(f"  全页解析后文档数: {len(dataprocess.data)}")

    print("\n📑 开始单页规则解析...")
    dataprocess.parse_one_page_with_rule(max_seq=256)
    dataprocess.parse_one_page_with_rule(max_seq=512)
    print(f"  单页解析后文档数: {len(dataprocess.data)}")

    # 构建 BM25 检索器
    print(f"\n✅ 构建 BM25 检索器，共加载 {len(dataprocess.data)} 个文档...")
    bm25 = BM25(dataprocess.data)

    # 执行查询
    query = "座椅加热"
    topk = 6
    print(f"\n🔎 正在检索: '{query}' (Top-{topk}) ...\n")
    print("=" * 80)

    results = bm25.get_bm25_topk(query, topk)

    # 清晰打印结果
    if not results:
        print("⚠️ 未找到相关文档。")
    else:
        for i, doc in enumerate(results, 1):
            metadata = doc.metadata
            content = doc.page_content.strip()
            print(f"🏆 第 {i} 名")
            print(f"📝 内容: {content}")
            print(f"🔖 元数据: ID={metadata['id']}")
            print("-" * 60)

    print("=" * 80)
    print(f"✅ 检索完成，共返回 {len(results)} 个结果。")
